# 4.6 使用分水岭和GrabCut算法进行物体分割

import cv2
import numpy
from matplotlib import pyplot as plt

if __name__ == "__main__" and False:

    img = cv2.imread('image/statue_small2.jpg')
    mask = numpy.zeros(img.shape[:2], numpy.uint8)

    bgd_model = numpy.zeros((1, 65), numpy.float64)
    fgd_model = numpy.zeros((1, 65), numpy.float64)

    rect = (100, 50, 421, 378)
    cv2.grabCut(img=img, mask=mask, rect=rect, bgdModel=bgd_model, fgdModel=fgd_model, iterCount=5, mode=cv2.GC_INIT_WITH_RECT)

    mask2 = numpy.where((mask == 2)|(mask == 0), 0, 1).astype('uint8')
    img = img * mask2[:, :, numpy.newaxis]

    plt.subplot(121), plt.imshow(img)
    plt.title("grabcut"), plt.xticks([]), plt.yticks([])
    plt.subplot(122), plt.imshow(cv2.cvtColor(cv2.imread("image/statue_small2.jpg"), cv2.COLOR_BGR2RGB))
    plt.title("original"), plt.xticks([]), plt.yticks([])
    plt.show()

    cv2.waitKey()
    cv2.destroyAllWindows()


# 使用分水岭算法进行图像分割
if __name__ == "__main__":

    img = cv2.imread('image/statue_small2.jpg')
    gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    ret, threshold = cv2.threshold(src=gray_img, thresh=0, maxval=255, type=cv2.THRESH_BINARY_INV+cv2.THRESH_OTSU)
    kernel = numpy.ones((3, 3), numpy.uint8)
    opening = cv2.morphologyEx(src=threshold, op=cv2.MORPH_OPEN, kernel=kernel, iterations=2)

    sure_bg = cv2.dilate(src=opening, kernel=kernel, iterations=3)

    dist_transform = cv2.distanceTransform(src=opening, distanceType=cv2.DIST_L2, maskSize=5)

    sure_fg = numpy.uint8(sure_bg)
    unknown = cv2.subtract(sure_bg, sure_fg)

    ret, markers = cv2.connectedComponents(image=sure_fg)

    markers = markers + 1
    markers[unknown == 255] = 0

    markers = cv2.watershed(image=img, markers=markers)
    img[markers == -1] = [0, 0, 0]
    plt.imshow(img)
    plt.show()

